#%% Imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import pingouin as pg
from pingouin import mixed_anova

#%% Import and subsetdata
df = pd.read_csv('/Users/juliafine/learnpy/combined_results.csv')
df.name = 'df'

par_ctrl = df.query('group != "control_new"')
par = df.query('group == "treatment"')

#%% Define a Mixed ANOVA function
def mix_anova (df, var_pre, var_post, value_name):

    #Get rid of nans in columns of interest
    df2 = df[~df[var_pre].isna()]
    df1 = df2[~df2[var_post].isna()]

    #Reorganize dataframe into long format
    df_melt = pd.melt(df1, id_vars=['ID', 'group'], value_vars=[var_pre, var_post],
                var_name='time', value_name=value_name)

    #Rename efficacy_prior and efficacy_post3 to "pre" and "post-3"
    df_melt['time'] = df_melt['time'].map({var_pre: 'pre', var_post: 'post-3'})

    aov = mixed_anova(dv=value_name, between='group',
                    within='time', subject='ID', data=df_melt)

    pg.print_table(aov)

    #Mauchly's test of sphericity
    print('Mauchly\'s test of sphericity')
    print(pg.sphericity(data=df_melt, dv=value_name, subject='ID', within='time')[-1])

    #Knowledge change - Shapiro-Wilks test that residuals are normally distributed
    print('Shapiro-Wilk\'s test that residuals are normally distributed')
    df_melt['factor_comb']=df_melt["group"] + '-' + df_melt["time"]
    print(pg.normality(df_melt, dv=value_name, group='factor_comb'))

    #Knowledge change - Levene's test of homoscedascicity
    df_melt_before = pd.melt(df1.reset_index(), id_vars=['ID', 'group'], value_vars=[var_pre])
    df_melt_after = pd.melt(df1.reset_index(), id_vars=['ID', 'group'], value_vars=[var_post])
    print('Levene\'s test of homoscedascicity - Pre-study')
    print(pg.homoscedasticity(df_melt_before, dv='value', group='group'))
    print('Levene\'s test of homoscedascicity - Post-3')
    print(pg.homoscedasticity(df_melt_after, dv='value', group='group'))

#%% Mixed anova -- efficacy change by group
mix_anova(par_ctrl, 'efficacy_prior', 'efficacy_3', 'efficacy_score')

#%% Mixed anova -- concern change by group
mix_anova(par_ctrl, 'concern_prior', 'concern_3', 'concern_score')

#%% Mixed anova -- likelihood change by group
mix_anova(par_ctrl, 'likelihood_prior', 'likelihood_3', 'likelihood_score')

#%% Mixed anova -- knowledge change by group
mix_anova(par_ctrl, 'knowledge_prior', 'knowledge_3', 'knowledge_score')

